import os
import numpy as np
import cv2
from PIL import Image

def letterbox_image(image, size):
    iw, ih = image.size
    w, h = size
    scale = min(w/iw, h/ih)
    nw = int(iw*scale)
    nh = int(ih*scale)

    image = image.resize((nw,nh), Image.BICUBIC)
    # new_image = Image.new('RGB', size, (128,128,128))
    new_image = Image.new('RGB', size, (0,0,0))
    new_image.paste(image, ((w-nw)//2, (h-nh)//2))
    return new_image

def search_files(directory):
    directory = os.path.normpath(directory)
    objects = {}
    for curdir, subdirs, files in os.walk(directory):
        for file in files:
            if file.endswith('.png'):
                label = curdir.split(os.path.sep)[-1]
                if label not in objects:
                    objects[label] = []
                path = os.path.join(curdir, file)
                objects[label].append(path)
    return objects
if __name__ == "__main__":
    train_samples = search_files('/root/image/car image/chest/seg')
    print(train_samples)
    for label, filenames in train_samples.items():
        for filename in filenames:
            img = Image.open(filename)
            # new_img = letterbox_image(img, (224, 224))
            new_img = letterbox_image(img, (512, 512))
            new_img.save(filename)
